"""
Load kissmetrics raw data from specified s3 bucket into specified mongodb.

Requires:
   * boto: Python library for AWS
   * pymongo: Python interface to mongo client

Comment:

This script can safely be run from an Amazon spot instance,
where checkpointing is important, since the script
can be interrupted at any moment.

Strategy:

kissmetrics deposits raw data into multiple json files:

    bucket/
        revisions/
            1.json
            2.json
            ...
            N.json

kissmetrics provides an 'index' file at:
    bucket/index.csv

This index file lists all the json files that kissmetrics
has currently deposited in the 'revisions' folder above.

Our strategy is to process the json files in the order
they appear in index.csv (which happens to be in order of
increasing time).

In our mongo instance, we will maintain our progress through
index.csv. This approach assumes that kissmetrics only ever
appends to index.csv, and never edits.
"""

import os
import csv
import json
try:
    import cStringIO as StringIO
except:
    import StringIO
import argparse

from boto.s3.connection import S3Connection
from pymongo import MongoClient

DEFAULT_INDEX_FILENAME = 'index.csv'

__loglevel = os.environ.get('ASA_LOGLEVEL', 'info')
def log(msg):
    is_debug = msg.startswith('DEBUG:')
    is_info = msg.startswith('INFO:')
    if (__loglevel == 'debug' and (is_info or is_debug)) or (__loglevel == 'info' and is_info):
        print msg

class km2mongoError(Exception):
    pass

def mkfile(bucket, keyname):
    f = StringIO.StringIO()
    key = bucket.get_key(keyname)
    key.get_contents_to_file(f)
    f.seek(0)
    return f

def gets3conn(aws_access_key, aws_secret_key):
    if aws_access_key and aws_secret_key:
        return S3Connection(aws_access_key, aws_secret_key)
    else:
        # look for env variables: AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY
        return S3Connection()

def get_current_progress(db, progcoll):
    """How many lines in index.csv have we processed?

    progcoll -> name of the mongo collection that maintains processing info.

    This collection maintains exactly one document, with one key-value pair:

    {'last_processed': 3455}

    If starting afresh, the value will be 0.
    """

    collection = db[progcoll]
    return collection.find_one()['last_processed']

def update_progress(db, progcoll, newlinenum):
    collection = db[progcoll]
    current_doc = collection.find_one()
    log("DEBUG: current_doc: %(current_doc)r" % locals())
    collection.update(current_doc, {"$set": {'last_processed': newlinenum }})

def get_new_data(bucket, indexfilename, current_progress):
    """Returns new json files that need to be processed.

    Return value is a list of tuples of the form:

    (linenum, jsonfile).

    The caller is responsible for updating progress
    upon successful handling of each json file.
    This is where linenum comes in handy.
    """
    f = mkfile(bucket, indexfilename)
    found = processed = 0
    reader = csv.reader(f, delimiter=',')
    for n, line in enumerate(reader, 0):
        # first line is header line: will also be skipped
        found += 1
        if n <= current_progress:
            continue
        jsonfile = line[0]  # first column is the json file location
        yield (n, jsonfile)
        processed += 1
    f.close()
    if not found:
        raise km2mongoError("kissmetrics index file is empty")
    if not processed:
        print "no new data to process!"
    else:
        print "found %(processed)d new files" % locals()

def _jsonload(f):
    for linenum, line in enumerate(f, 1):
        try:
            yield json.loads(line)
        except:
            log("DEBUG: unable to parse json data from line num %(linenum)s: %(line)s" % locals())

def jsonfile_to_mongo(bucket, linenum, jsonfile, db, progcoll, datacoll):
    collection = db[datacoll]
    f = mkfile(bucket, jsonfile)
    data = [doc for doc in _jsonload(f)]
    # the following should happen in a transaction;
    # but mongo doesn't support this? what is this shit.
    collection.insert(data)
    update_progress(db, progcoll, linenum)

def km2mongo(s3bucketname,
             mongohost, mongoport, mongodbname, mongo_username, mongo_passwd,
             mongoprogcoll, mongodatacoll,
             indexfilename=DEFAULT_INDEX_FILENAME,
             aws_access_key=None, aws_secret_key=None):
    # --- 1. set up connection to mongo
    db = MongoClient(mongohost, mongoport)[mongodbname]
    db.authenticate(mongo_username, password=mongo_passwd)
    # --- 2. set up connection to aws
    s3conn = gets3conn(aws_access_key, aws_secret_key)
    bucket = s3conn.get_bucket(s3bucketname)
    # --- 3. actual work
    p = get_current_progress(db, mongoprogcoll)
    log("DEBUG: current progress: %(p)s" % locals())
    for linenum, jsonfile in get_new_data(bucket, indexfilename, p):
        log("DEBUG: processing line %(linenum)s from file %(indexfilename)s: %(jsonfile)s ..." % locals())
        jsonfile_to_mongo(bucket, linenum, jsonfile, db, mongoprogcoll, mongodatacoll)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="load kissmetrics data from specified s3 bucket into specified mongodb")
    parser.add_argument('-b', '--s3bucket', help="name of s3 bucket where kissmetrics stores data", required=True)
    parser.add_argument('-m', '--mongohost', help="hostname of machine hosting mongo db instance", required=True)
    parser.add_argument('-p', '--mongoport', help="port on which mongo db instance listens", type=int, required=True)
    parser.add_argument('-d', '--mongodbname', help="name of mongo database", required=True)
    parser.add_argument('-c', '--datacollection', help="name of mongodb collection on which to store actual kissmetrics data", required=True)
    parser.add_argument('-x', '--proccollection', help="name of mongodb collection on which to store progress data", required=True)
    parser.add_argument('-u', '--user', help="mongodb username", required=True)
    parser.add_argument('-w', '--passwd', help="mongodb password", required=True)
    parser.add_argument('-i', '--indexfilename', help="name of kissmetrics index file", default=DEFAULT_INDEX_FILENAME)
    parser.add_argument('-a', '--aws-access-key', help="aws access key; if not specified, "
                        "script will look for following env var: AWS_ACCESS_KEY_ID")
    parser.add_argument('-s', '--aws-secret-key', help="aws secret key; if not specified, "
                        "script will look for following env var: AWS_SECRET_ACCESS_KEY")
    args = parser.parse_args()
    km2mongo(args.s3bucket,
             args.mongohost, args.mongoport, args.mongodbname, args.user, args.passwd,
             args.proccollection, args.datacollection,
             indexfilename=args.indexfilename,
             aws_access_key=args.aws_access_key, aws_secret_key=args.aws_secret_key)
